# -*- coding: UTF-8 -*-
'''
@Project ：push_rk
@File ：model.py
@IDE ：PyCharm
@Author ：苦瓜
@Date ：2025/9/24 10:35
@Note: Something beautiful is about to happen !
'''
import numpy as np
from keras.src.datasets import cifar100
from keras import utils, Model, Sequential, optimizers, losses, metrics
from keras.src.layers import Conv2D, Flatten, Dense, Dropout, MaxPooling2D

class AlexNet(Model):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 卷积层
        self.conv = Sequential([
            Conv2D(filters=6, kernel_size=3, activation="relu", padding="same"),
            MaxPooling2D(),
            Conv2D(filters=16, kernel_size=3, activation="relu", padding="same"),
            MaxPooling2D(),
            Conv2D(filters=24, kernel_size=3, activation="relu", padding="same"),
            Conv2D(filters=24, kernel_size=3, activation="relu", padding="same"),
            Conv2D(filters=16, kernel_size=3, activation="relu", padding="same"),
            MaxPooling2D()
        ])
        # 展平层
        self.flatten = Flatten()
        # 全连接层
        self.dense = Sequential([
            Dense(units=64, activation='relu'),
            Dense(units=64, activation='relu'),
            Dense(units=100, activation='softmax')
        ])

    def call(self, inputs, *args, **kwargs):
        x = self.conv(inputs)
        x = self.flatten(x)
        x = self.dense(x)
        return x


class MyBaseCNN:
    _product_dict = {
    0: 'apple',
    1: 'aquarium_fish',
    2: 'baby',
    3: 'bear',
    4: 'beaver',
    5: 'bed',
    6: 'bee',
    7: 'beetle',
    8: 'bicycle',
    9: 'bottle',
    10: 'bowl',
    11: 'boy',
    12: 'bridge',
    13: 'bus',
    14: 'butterfly',
    15: 'camel',
    16: 'can',
    17: 'castle',
    18: 'caterpillar',
    19: 'cattle',
    20: 'chair',
    21: 'chimpanzee',
    22: 'clock',
    23: 'cloud',
    24: 'cockroach',
    25: 'couch',
    26: 'crab',
    27: 'crocodile',
    28: 'cup',
    29: 'dinosaur',
    30: 'dolphin',
    31: 'elephant',
    32: 'flatfish',
    33: 'forest',
    34: 'fox',
    35: 'girl',
    36: 'hamster',
    37: 'house',
    38: 'kangaroo',
    39: 'keyboard',
    40: 'lamp',
    41: 'lawn_mower',
    42: 'leopard',
    43: 'lion',
    44: 'lizard',
    45: 'lobster',
    46: 'man',
    47: 'maple_tree',
    48: 'motorcycle',
    49: 'mountain',
    50: 'mouse',
    51: 'mushroom',
    52: 'oak_tree',
    53: 'orange',
    54: 'orchid',
    55: 'otter',
    56: 'palm_tree',
    57: 'pear',
    58: 'pickup_truck',
    59: 'pine_tree',
    60: 'plain',
    61: 'plate',
    62: 'poppy',
    63: 'porcupine',
    64: 'possum',
    65: 'rabbit',
    66: 'raccoon',
    67: 'ray',
    68: 'road',
    69: 'rocket',
    70: 'rose',
    71: 'sea',
    72: 'seal',
    73: 'shark',
    74: 'shrew',
    75: 'skunk',
    76: 'skyscraper',
    77: 'snail',
    78: 'snake',
    79: 'spider',
    80: 'squirrel',
    81: 'streetcar',
    82: 'sunflower',
    83: 'sweet_pepper',
    84: 'table',
    85: 'tank',
    86: 'telephone',
    87: 'television',
    88: 'tiger',
    89: 'tractor',
    90: 'train',
    91: 'trout',
    92: 'tulip',
    93: 'turtle',
    94: 'wardrobe',
    95: 'whale',
    96: 'willow_tree',
    97: 'wolf',
    98: 'woman',
    99: 'worm'
}  # 预测百类别得字典

    def __init__(self):
        self.model = None
        self.history = None

    def load_data(self, datasets_module):
        """ 加载数据 """
        (X_train, y_train), (X_test, y_test) = datasets_module.load_data()  # 获取数据
        # 修改数据维度
        X_train, X_test = X_train.reshape(-1, 32, 32, 3), X_test.reshape(-1, 32, 32, 3)
        # 归一化
        X_train, X_test = X_train.astype('float32') / 255. , X_test.astype('float32') / 255.
        # 获取类别数量
        y_dim = len(np.unique(y_train))
        # 对标签进行独热
        y_train, y_test = self.one_hot_encoding(y_train, y_dim), self.one_hot_encoding(y_test, y_dim)
        return X_train, X_test, y_train, y_test

    def one_hot_encoding(self, x, cls_num):
        """ 对目标数据进行独热 """
        return utils.to_categorical(x, cls_num)

    def model_build(self, X_train, y_train, batch_size=100, epochs=10, validation_data=None):
        """ 构建模型、编译、训练 """
        self.model = AlexNet()
        self.model.build(input_shape=(None, 32, 32, 3))
        # 模型编译
        self.model.compile(
            optimizer=optimizers.Adam(),
            loss=losses.categorical_crossentropy,
            metrics=['acc', metrics.Recall, metrics.Precision, metrics.F1Score]
        )
        # 模型训练
        self.history = self.model.fit(
            X_train, y_train, batch_size=batch_size,
            epochs=epochs, validation_data=validation_data
        )
        return self.model, self.history
    @classmethod
    def predict_cls(cls, num):
        return cls._product_dict[num]

def main():
    # 获取基础模型
    base_obj = MyBaseCNN()
    # 获取数据
    X_train, X_test, y_train, y_test = base_obj.load_data(cifar100)
    # 构建模型并编译训练
    model, history = base_obj.model_build(X_train,y_train,
                                          validation_data=(X_test, y_test)
                                          )
    # 打印模型结构
    model.summary()
    # 模型评估
    print(model.evaluate(X_test, y_test))
    # 保存模型得权重
    model.save_weights("model.weights.h5")

if __name__ == '__main__':
    main()  # 逻辑入口


# 1)	数据加载与预处理：
# 加载CIFAR-100数据集，并将其分为训练集和测试集。
# 对图像的像素值进行归一化处理，将像素值从[0, 255]缩放到[0, 1]，以改善模型的训练效果。
# 2)	模型定义：
# 定义一个名为AlexNet的自定义模型类。
# 在模型类中，构建了一个卷积神经网络架构，该架构包括五个卷积层（每个卷积层后都有一个最大池化层），用于提取图像中的特征。
# 接着是一个展平层，用于将卷积层输出的多维张量转换为一维数组，以便输入到全连接层中。
# 最后是两个全连接层（每个全连接层后都有一个ReLU激活函数）和一个输出层（具有100个神经元和softmax激活函数），用于对图像进行分类。
# 3)	模型编译：
# 使用合适的优化器来编译模型，优化器是一种自适应学习率的方法，能够有效地训练深度神经网络。
# 损失函数设置为稀疏分类交叉熵，因为CIFAR-100的标签是整数形式的。
# 评估指标设置为准确率。
# 4)	模型训练：
# 使用训练集x_train和y_train对模型进行训练，批处理大小为100，训练轮次为10。
# 5)	模型评估：
# 使用测试集x_test和y_test对训练好的模型进行评估，计算模型在测试集上的准确率。
# 6)	权重保存：
# 将训练好的模型的权重保存为HDF5格式的文件（'model.h5'），以便后续使用或进一步分析。
